import os
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from FSdata.FSdataset import FSdata, collate_fn,attr2length_map, idx2attr_map
import torch
import torch.utils.data as torchdata
from models.resnet import resnet50_cat,resnet101_cat_dilate
from FSdata.FSaug import *
from utils.predicting import predict,gen_submission
from utils.metrics import cal_mAP
from models.NASnet import NASnet_cat
from utils.preprocessing import join_path_to_df

class FSAugTest(object):
    def __init__(self,doHflip,ratio=1,crop=2,nochange=0):
        self.augment = Compose([
            UpperCrop(size=(400, 400), select=[1,2, 3, 4],ratio=ratio,crop=crop,nochange=nochange),
            # ResizedCrop(size=(336, 336), select=[ 7],ratio=ratio,crop=crop,nochange=nochange),
            # DownCrop(size=(336, 336), select=[5, 6],ratio=ratio,crop=crop,nochange=nochange),
            Resize(size=(400, 400), select=[0, 1, 2, 3, 4, 5, 6, 7]),
            Hflip(doHflip=doHflip),

            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, image, attr):
        return self.augment(image, attr)


os.environ["CUDA_VISIBLE_DEVICES"] = "1"

rawdata_root = '/media/gserver/data/FashionAI'
round2_df = pd.read_csv(os.path.join(rawdata_root,'round2/train/Annotations/label.csv'),
                        header=None, names=['ImageName', 'AttrKey', 'AttrValues'])
round2_df = join_path_to_df(round2_df, rawdata_root, 'round2/train')

round2_train_pd, val_pd = train_test_split(round2_df, test_size=0.1, random_state=37,
                                    stratify=round2_df['AttrKey'])


# select part
# select_AttrIdx = range(8)
select_AttrIdx = [4]


select_AttrKey = [idx2attr_map[x] for x in select_AttrIdx]
val_pd = val_pd[val_pd['AttrKey'].apply(lambda x: True if x in select_AttrKey else False)]


print val_pd.shape
print val_pd['AttrKey'].value_counts()


# model prepare

# model prepare
model_name = 'res101_d[r2-4]9748-aug14'
resume = '/media/gserver/extra/FashionAI/round2/res101_d[1234]_seed0_[r2-4]/weights-15-1000-[0.8852]-[0.9748].pth'

model = resnet101_cat_dilate(pretrained=True, num_classes=[attr2length_map[x] for x in select_AttrIdx])

# model = torch.nn.DataParallel(model)
print('resuming finetune from %s'%resume)
model.load_state_dict(torch.load(resume))
# model = model.module
model = model.cuda()


# predict
if not os.path.exists('./val_pred2/csv'):
    os.makedirs('./val_pred2/csv')
if not os.path.exists('./val_pred2/part_csv'):
    os.makedirs('./val_pred2/part_csv')


aug_scores = np.zeros((14, val_pd.shape[0], sum([attr2length_map[x] for x in select_AttrIdx])), dtype=np.float32)


for i,doHflip in enumerate([0,1]):
    for j,ratio in enumerate([7. /8. , 7./6.]):
        for k,crop in enumerate([0,1,2]):

            data_set = {}
            data_set['val'] = FSdata( anno_pd=val_pd,
                                      transforms=FSAugTest(doHflip=doHflip,ratio=ratio,crop=crop),
                                      select=select_AttrIdx
                                      )

            data_loader = {}
            data_loader['val'] = torchdata.DataLoader(data_set['val'], batch_size=16, num_workers=8,
                                                      shuffle=False, pin_memory=True,collate_fn=collate_fn)

            aug_scores[i*7+j*3+k,:,:],val_labels_str, val_attr, _ = predict(model, data_set['val'], data_loader['val'])
    data_set = {}
    data_set['val'] = FSdata(anno_pd=val_pd,
                             transforms=FSAugTest(doHflip=doHflip, ratio=1, crop=1,nochange=1),
                             select=select_AttrIdx
                             )

    data_loader = {}
    data_loader['val'] = torchdata.DataLoader(data_set['val'], batch_size=16, num_workers=8,
                                              shuffle=False, pin_memory=True, collate_fn=collate_fn)
    aug_scores[i * 7 + 6, :, :], val_labels_str, val_attr, _ = predict(model, data_set['val'],
                                                                               data_loader['val'])

pred_scores = aug_scores.mean(0)

val_mAP,APs, accs = cal_mAP(val_labels_str,pred_scores,val_attr,data_set['val'].catidx_map)

print('=='*20)
print('val-mAP: %.4f'% (val_mAP))
for key in APs.keys():
    print('acc: %.4f, AP: %.4f %s' % (accs[key], APs[key], key))


val_pred = gen_submission(val_pd[['ImageName', 'AttrKey']], pred_scores, catidx_map=data_set['val'].catidx_map)

if len(select_AttrIdx) < 8:
    val_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('./val_pred2/part_csv/%s-[%.4f].csv' % (model_name,val_mAP),
                                                                 header=None, index=False)
else:
    val_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('./val_pred2/csv/%s-[%.4f].csv' % (model_name,val_mAP),
                                                                 header=None, index=False)

print val_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].info()